// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include "column/vectorized_fwd.h"
#include "exec/pipeline/operator.h"
#include "exprs/expr.h"
#include "exprs/table_function/table_function_factory.h"
#include "runtime/descriptors.h"
#include "runtime/runtime_state.h"

namespace starrocks::pipeline {
class TableFunctionOperator final : public Operator {
public:
    TableFunctionOperator(OperatorFactory* factory, int32_t id, int32_t plan_node_id, int32_t driver_sequence,
                          const TPlanNode& tnode)
            : Operator(factory, id, "table_function", plan_node_id, false, driver_sequence), _tnode(tnode) {}

    ~TableFunctionOperator() override = default;

    [[nodiscard]] Status prepare(RuntimeState* state) override;

    void close(RuntimeState* state) override;

    bool has_output() const override;

    bool need_input() const override;

    bool is_finished() const override;

    [[nodiscard]] Status set_finishing(RuntimeState* state) override;

    StatusOr<ChunkPtr> pull_chunk(RuntimeState* state) override;

    [[nodiscard]] Status push_chunk(RuntimeState* state, const ChunkPtr& chunk) override;

    [[nodiscard]] Status reset_state(RuntimeState* state, const std::vector<ChunkPtr>& refill_chunks) override;

private:
    ChunkPtr _build_chunk(const std::vector<ColumnPtr>& output_columns);
    [[nodiscard]] Status _process_table_function(RuntimeState* state);
    void _copy_result(const std::vector<ColumnPtr>& columns, uint32_t max_column_size);

    const TPlanNode& _tnode;
    const TableFunction* _table_function = nullptr;

    // Slots of output by table function
    std::vector<SlotId> _fn_result_slots;
    // External column slots of the join logic generated by the table function
    std::vector<SlotId> _outer_slots;
    // Slots of table function input parameters
    std::vector<SlotId> _param_slots;

    // Input chunk currently being processed
    ChunkPtr _input_chunk;
    // The subscript in "input_chunk" of the input row corresponding to the first row of "_table_function_result.first".
    size_t _input_index_of_first_result = 0;
    // How many rows of "_table_function_result.first" have been output to downstream operator.
    size_t _next_output_row = 0;
    // The subscript in "_table_function_result.second" corresponding to the "_next_output_row".
    size_t _next_output_row_offset = 0;
    // table function result
    std::pair<Columns, UInt32Column::Ptr> _table_function_result;
    // table function param and return offset
    TableFunctionState* _table_function_state = nullptr;

    RuntimeProfile::Counter* _table_function_exec_timer = nullptr;
    RuntimeProfile::Counter* _table_function_exec_counter = nullptr;

    bool _is_finished = false;
};

class TableFunctionOperatorFactory final : public OperatorFactory {
public:
    TableFunctionOperatorFactory(int32_t id, int32_t plan_node_id, const TPlanNode& tnode)
            : OperatorFactory(id, "table_function", plan_node_id), _tnode(tnode) {}

    OperatorPtr create(int32_t degree_of_parallelism, int32_t driver_sequence) override {
        return std::make_shared<TableFunctionOperator>(this, _id, _plan_node_id, driver_sequence, _tnode);
    }

private:
    const TPlanNode& _tnode;
};

} // namespace starrocks::pipeline
